from abc import abstractmethod
from typing import Any, Dict, List, Optional
import torch
from torch import nn


class TaskAwareInterface:
    """Interface that are aware of multiple tasks."""

    _current_task: int = -1

    @abstractmethod
    def register_task(self, task_id: int, **kwargs) -> None:
        """
        Register a new task with this hook.

        Args:
            task_id: Identifier for the new task
            config: Optional task-specific configuration
        """
        pass

    @abstractmethod
    def set_task(self, task_id: int) -> None:
        """
        Configure the model for a specific task.

        Args:
            task_id: Task identifier
        """
        self._current_task = task_id

    @abstractmethod
    def set_trainable_tasks(self, task_ids: int | List[int], **kwargs) -> None:
        """
        Prepare this hook for a specific task.

        Args:
            task_id: Identifier for the task to prepare
            **kwargs: Additional preparation parameters
        """
        pass

    @abstractmethod
    def _freeze(self, task_ids: int | List[int]) -> None:
        """
        Freeze parameters for a specific task.

        Args:
            task_id: Task identifier to freeze parameters for
        """
        pass

    @abstractmethod
    def _unfreeze(self, task_ids: int | List[int]) -> None:
        """
        Freeze parameters for a specific task.

        Args:
            task_id: Task identifier to freeze parameters for
        """
        pass

    @abstractmethod
    def process_with_task(
        self, x: torch.Tensor, task_id: int, **kwargs
    ) -> torch.Tensor:
        """Process input with a specific task."""
        pass

    @abstractmethod
    def process_with_tasks(
        self, x: torch.Tensor, task_ids: Optional[List[int]] = None, **kwargs
    ) -> torch.Tensor:
        """Process input among multiple tasks."""
        pass

    @property
    @abstractmethod
    def current_task(self) -> int:
        """Get the current task ID."""
        pass

    @property
    @abstractmethod
    def registered_tasks(self) -> List[int]:
        """Get list of registered task IDs."""
        pass


class TaskIrrelevantInterface:
    """Interface that are not aware of multiple tasks."""

    @abstractmethod
    def get_parameters(self) -> List[nn.Parameter]:
        """
        Get trainable parameters managed by this module.

        Returns:
            List[nn.Parameter]: List of parameters
        """
        pass

    @abstractmethod
    def get_trainable_parameters(self) -> List[nn.Parameter]:
        """
        Get trainable parameters managed.

        Returns:
            List[nn.Parameter]: List of parameters
        """
        pass
